import torch
from torch import nn
import torch.nn.functional as F
import torch.optim as optim


def contracting_block(in_channels, out_channels, kernel_size=3):
    block = torch.nn.Sequential(
        torch.nn.Conv2d(kernel_size=kernel_size, in_channels=in_channels, out_channels=out_channels, padding=1),
        torch.nn.ReLU(),
        torch.nn.BatchNorm2d(out_channels),
        torch.nn.Conv2d(kernel_size=kernel_size, in_channels=out_channels, out_channels=out_channels, padding=1),
        torch.nn.ReLU(),
        torch.nn.BatchNorm2d(out_channels),
    )
    return block


class Attention_block(nn.Module):
    def __init__(self, F_g, F_l, F_int):
        '''

        :param F_g:  upsample
        :param F_l: skip
        :param F_int:
        '''
        super(Attention_block, self).__init__()
        self.W_g = nn.Sequential(
            nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(F_int)
        )

        self.W_x = nn.Sequential(
            nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(F_int)
        )

        self.psi = nn.Sequential(
            nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True),
            nn.BatchNorm2d(1),
            nn.Sigmoid()
        )

        self.relu = nn.ReLU(inplace=True)

    def forward(self, g, x):
        g1 = self.W_g(g)
        x1 = self.W_x(x)

        c = (g1.size()[2] - x1.size()[2]) // 2
        w = g1.size()[2] - x1.size()[2]
        h = g1.size()[2] - x1.size()[3]
        if w % 2 == 0 and h % 2 == 0:
            # bypass = F.pad(bypass, (-c, -c, -c, -c))
            g1 = F.pad(g1, (c, c, c, c))
        else:
            g1 = F.pad(g1, (1, 0, 1, 0))


        psi = self.relu(g1 + x1)
        psi = self.psi(psi)

        return x * psi



class Unet(nn.Module):

    def expansive_block(self, in_channels, mid_channel, out_channels, kernel_size=3):
        block = torch.nn.Sequential(
            torch.nn.Conv2d(kernel_size=kernel_size, in_channels=in_channels, out_channels=mid_channel, padding=1),
            torch.nn.ReLU(),
            torch.nn.BatchNorm2d(mid_channel),
            torch.nn.Conv2d(kernel_size=kernel_size, in_channels=mid_channel, out_channels=mid_channel, padding=1),
            torch.nn.ReLU(),
            torch.nn.BatchNorm2d(mid_channel),
            torch.nn.ConvTranspose2d(in_channels=mid_channel, out_channels=out_channels, kernel_size=3, stride=2,
                                     padding=1, output_padding=1)
        )
        return block

    def final_block(self, in_channels, mid_channel, out_channels, kernel_size=3):
        block = torch.nn.Sequential(
            torch.nn.Conv2d(kernel_size=kernel_size, in_channels=in_channels, out_channels=mid_channel, padding=1),
            torch.nn.ReLU(),
            torch.nn.BatchNorm2d(mid_channel),
            torch.nn.Conv2d(kernel_size=kernel_size, in_channels=mid_channel, out_channels=mid_channel, padding=1),
            torch.nn.ReLU(),
            torch.nn.BatchNorm2d(mid_channel),
            torch.nn.Conv2d(kernel_size=kernel_size, in_channels=mid_channel, out_channels=out_channels, padding=1),
            torch.nn.ReLU(),
            torch.nn.BatchNorm2d(out_channels),
        )
        return block


    def __init__(self, in_channel=3, out_channel=1):
        super(Unet, self).__init__()
        # Encode
        self.conv_encode1 = contracting_block(in_channels=in_channel, out_channels=64)
        self.conv_maxpool1 = torch.nn.MaxPool2d(kernel_size=2)
        self.conv_encode2 = contracting_block(64, 128)
        self.conv_maxpool2 = torch.nn.MaxPool2d(kernel_size=2)
        self.conv_encode3 = contracting_block(128, 256)
        self.conv_maxpool3 = torch.nn.MaxPool2d(kernel_size=2)
        self.conv_encode4 = contracting_block(256, 512)
        self.conv_maxpool4 = torch.nn.MaxPool2d(kernel_size=2)
        # Bottleneck
        self.bottleneck = torch.nn.Sequential(
            torch.nn.Conv2d(kernel_size=3, in_channels=512, out_channels=1024, padding=1),
            torch.nn.ReLU(),
            torch.nn.BatchNorm2d(1024),
            torch.nn.Conv2d(kernel_size=3, in_channels=1024, out_channels=1024, padding=1),
            torch.nn.ReLU(),
            torch.nn.BatchNorm2d(1024),
            torch.nn.ConvTranspose2d(in_channels=1024, out_channels=512, kernel_size=3, stride=2, padding=1,
                                     output_padding=1)
        )
        # Attention
        self.Att4 = Attention_block(F_g=512, F_l=512, F_int=512)
        self.Att3 = Attention_block(F_g=256, F_l=256, F_int=256)
        self.Att2 = Attention_block(F_g=128, F_l=128, F_int=128)
        self.Att1 = Attention_block(F_g=64, F_l=64, F_int=64)
        # Decode
        self.conv_decode4 = self.expansive_block(1024, 512, 256)
        self.conv_decode3 = self.expansive_block(512, 256, 128)
        self.conv_decode2 = self.expansive_block(256, 128, 64)
        self.final_layer = self.final_block(128, 64, out_channel)



    def crop_and_concat(self, upsampled, bypass, crop=False):
        if crop:
            c = (bypass.size()[2] - upsampled.size()[2]) // 2
            w = bypass.size()[2] - upsampled.size()[2]
            h = bypass.size()[2] - upsampled.size()[3]
            if w % 2 == 0 and h % 2 == 0:
                # bypass = F.pad(bypass, (-c, -c, -c, -c))
                upsampled = F.pad(upsampled, (c, c, c, c))
            else:
                upsampled = F.pad(upsampled, (1, 0, 1, 0))
        return torch.cat((upsampled, bypass), 1)

    def forward(self, x):
        # Encode
        encode_block1 = self.conv_encode1(x)
        encode_pool1 = self.conv_maxpool1(encode_block1)
        encode_block2 = self.conv_encode2(encode_pool1)
        encode_pool2 = self.conv_maxpool2(encode_block2)
        encode_block3 = self.conv_encode3(encode_pool2)
        encode_pool3 = self.conv_maxpool3(encode_block3)
        encode_block4 = self.conv_encode4(encode_pool3)
        encode_pool4 = self.conv_maxpool4(encode_block4)
        # Bottleneck
        bottleneck1 = self.bottleneck(encode_pool4)
        # Decode
        attention_block4 = self.Att4(bottleneck1, encode_block4)
        decode_block4 = self.crop_and_concat(bottleneck1, attention_block4, crop=True)
        cat_layer3 = self.conv_decode4(decode_block4)
        attention_block3 = self.Att3(cat_layer3, encode_block3)
        decode_block3 = self.crop_and_concat(cat_layer3, attention_block3, crop=True)
        cat_layer2 = self.conv_decode3(decode_block3)
        attention_block2 = self.Att2(cat_layer2, encode_block2)
        decode_block2 = self.crop_and_concat(cat_layer2, attention_block2, crop=True)
        cat_layer1 = self.conv_decode2(decode_block2)
        attention_block1 = self.Att1(cat_layer1, encode_block1)
        decode_block1 = self.crop_and_concat(cat_layer1, attention_block1, crop=True)
        final_layer = self.final_layer(decode_block1)
        return F.sigmoid(final_layer)
